################################################################################
## This script generates Figures 1, 2, and 3.
## Images are saved to the directory provided as a command line argument.

## USAGE: Rscript generate_figures.R ${OUTPUT_SAVE_DIR}
################################################################################


args = commandArgs(trailingOnly=TRUE)
if (length(args)!=1) {
  stop("ERROR: Missing argument: Include a directory in which to save the figures.")
}

SAVEDIR <- as.character(args[1])

kFigureWidth <- 5

############# FIGURE 1 :: Section 3.1 (simulation)

pdf(file=file.path(SAVEDIR, sprintf("fig1_intution.pdf")), height=kFigureWidth, width=kFigureWidth)
  par(pty="s", family = 'serif')
  tt=0.1
  curve((tt-x)/(1-x)*1*((tt-x)/(1-x)>0),0,1, xlim=c(0, 1), ylim=c(0,1), ylab="w0+w1*x",lty=3)
  curve(tt/(1-x)*1*(tt/(1-x)<1)+1*(tt/(1-x)>1),add=TRUE,lty=3)
  curve(tt+2*x,add=TRUE,lty=2)
  curve(tt+0.9*x,add=TRUE,lty=2)
  curve(tt+0.1*x,add=TRUE,lty=1)
  curve(tt-0.1*x,add=TRUE,lty=1)
  text(0.90,0.7,"E(U|X)")
  text(0.3,-0.0,"E(L|X)")
  text(0.2,0.7,"0.1+2X")
  text(0.5,0.7,"0.1+0.9X")
  text(0.7,0.1,"0.1+/-0.1X")
dev.off()


############# Figure 2 :: Section 6, Example 4 (simulation)

library(sandwich)
set.seed(7)

pdf(file=file.path(SAVEDIR, sprintf("fig2_example4.pdf")), height=kFigureWidth, width=kFigureWidth)
  par(pty="s", family = 'serif')
  p=1000 # number of precincts, in practice should be observed
  nn=150 
  n=rep(nn,p) #precinct sizes, in practice should be observed
  x=runif(p,0,0.95) #proportions of blacks, in practice should be observed
  l=min(x)  
  u=max(x) 
  ww0=0.9 #white vote prob in pure white precinct, unavailable in practice
  ww1=0.9 #white vote prob in pure black precinct, unavailable in practice
  bb0=0.9 #black vote prob in pure white precinct, unavailable in practice
  bb1=0.6 #black vote prob in pure black precinct, unavailable in practice
  w0=log(ww0/(1-ww0))
  #in practice w0 should be estimated by regressing t<-w0+c1x+d1x^2
  w1=log(ww1/(1-ww1))-log(ww0/(1-ww0))
  b0=log(bb0/(1-bb0))
  b1=log(bb1/(1-bb1))-log(bb0/(1-bb0))
  b=rep(0.5,p)
  w=b
  s=0.5 #extra noise sd
  for(i in 1:p) {
    b[i]=rbinom(1,round(n[i]*x[i])+1,1/(1+exp(-b0-b1*x[i]+s*(1-x[i])*rnorm(1))))/(round(n[i]*x[i])+1)
    w[i]=rbinom(1,round(n[i]*(1-x[i]))+1,1/(1+exp(-w0-w1*x[i]+s*(1-x[i])*rnorm(1))))/(round(n[i]*(1-x[i]))+1)
  }
  t=(w*round(n*(1-x))+b*round(n*x))/n #comb'd vote prop, in practice is observed
  plot(x,t, xlim=c(0, 1), ylim=c(0,1))
dev.off()


############# Figure 3 :: Section 6, Example 5 (simulation)

library(sandwich)
set.seed(7)

pdf(file=file.path(SAVEDIR, sprintf("fig3_example5.pdf")), height=kFigureWidth, width=kFigureWidth)
  par(pty="s", family = 'serif')
  p=1000 # number of precincts, in practice should be observed
  nn=150 
  n=rep(nn,p) #precinct sizes, in practice should be observed
  x=runif(p,0,0.7) #proportions of blacks, in practice should be observed
  l=min(x) 
  u=max(x) 
  ww0=0.9 #white vote prob in pure white precinct, unavailable in practice
  ww1=0.9 #white vote prob in pure black precinct, unavailable in practice
  bb0=0.5 #black vote prob in pure white precinct, unavailable in practice
  bb1=0.5 #black vote prob in pure black precinct, unavailable in practice
  w0=log(ww0/(1-ww0))
  #in practice w0 should be estimated by regressing t<-w0+c1x+d1x^2
  w1=log(ww1/(1-ww1))-log(ww0/(1-ww0))
  b0=log(bb0/(1-bb0))
  b1=log(bb1/(1-bb1))-log(bb0/(1-bb0))
  b=rep(0.5,p)
  w=b
  s=1 #extra noise sd
  for(i in 1:p) {
    b[i]=rbinom(1,round(n[i]*x[i])+1,1/(1+exp(-b0-b1*x[i]+s*rnorm(1))))/(round(n[i]*x[i])+1)
    w[i]=rbinom(1,round(n[i]*(1-x[i]))+1,1/(1+exp(-w0-w1*x[i]+s*rnorm(1))))/(round(n[i]*(1-x[i]))+1)
  }
  t=(w*round(n*(1-x))+b*round(n*x))/n #comb'd vote prop, in practice is observed
  plot(x,t, xlim=c(0, 1), ylim=c(0,1))
dev.off()
